from dataclasses import dataclass, field
from typing import Any, Tuple, Dict
import torch.nn as nn

@dataclass
class Config:
    env_id = "ale/icehockey-v5"
    pixel: bool = True
    gamma = 0.99
    batch_size = 16
    chunk_size = 50
    model_learning_rate = 2e-4
    actor_learning_rate = 4e-5
    value_learning_rate = 1e-5
    loss_scale: Dict = field(default_factory=lambda:{'kl':0.1, 'reward':1.0, 'discount':5.0})
    kl: Dict = field(default_factory=lambda:{'use_kl_balance':True, 'kl_balance_scale':0.8, 'use_free_nats':False, 'free_nats':0.0})
    adam_epsilon = 1e-7
    learning_rate_schedule = 0
    replay_size = 100000 # 重放缓冲区长度，这么长是为了提高稳定性
    replay_initial = 10000 # 重放缓冲区初始化大小
    collect_intervals: int = 5 
    lambda_: float = 0.95

    test_iters = 10

    precision = 32
    state_size = 30
    hide_size = 200
    belief_size = 200
    reward_step = 1
    embedding_size = 128
    free_nats = float(3.0)
    overshooting_distance = 50

    overshooting_reward_scale = float(0)
    planning_horizon = 10
    grad_clip_norm = 100
    train_every = 5
    slow_target_update: int = 50

    disclam = 0.95
    action_noise = 0.3

    action_size = 0
    obs_shape = 0

    rssm_node_size = 200

    rssm_type: str = 'discrete'
    rssm_info: Dict = field(default_factory=lambda:{'deter_size':100, 'stoch_size':256, 'class_size':16, 'category_size':16, 'min_std':0.1})

    actor: Dict = field(default_factory=lambda:{'layers':3, 'node_size':100, 'dist':'one_hot', 'min_std':1e-4, 'init_std':5, 'mean_scale':5, 'activation':nn.ELU})
    critic: Dict = field(default_factory=lambda:{'layers':3, 'node_size':100, 'dist': 'normal', 'activation':nn.ELU})
    expl: Dict = field(default_factory=lambda:{'train_noise':0.4, 'eval_noise':0.0, 'expl_min':0.05, 'expl_decay':7000.0, 'expl_type':'epsilon_greedy'})
    actor_grad: str ='reinforce'
    actor_grad_mix: int = 0.0
    actor_entropy_scale: float = 1e-3

    modelstate_size = 0

    obs_encoder: Dict = field(default_factory=lambda:{'layers':3, 'node_size':100, 'dist': None, 'activation':nn.ELU, 'kernel':3, 'depth':16})
    obs_decoder: Dict = field(default_factory=lambda:{'layers':3, 'node_size':100, 'dist':'normal', 'activation':nn.ELU, 'kernel':3, 'depth':16})
    reward: Dict = field(default_factory=lambda:{'layers':3, 'node_size':100, 'dist':'normal', 'activation':nn.ELU})
    discount: Dict = field(default_factory=lambda:{'layers':3, 'node_size':100, 'dist':'binary', 'activation':nn.ELU, 'use':True})
    discount_: float = 0.99


